#include "utils.h"
#include <iostream>
#include <string>
#include <mutex>
#include <chrono>
#include <boost/asio.hpp>

void print_banner() {
    std::cout << "================================================" << std::endl;
    std::cout << "  Port Scanner - Scan TCP ports on specified IPs" << std::endl;
    std::cout << "================================================" << std::endl;
}

bool is_port_open(const std::string& ip, int port) {
    // 移除调试输出
    // std::cout << "Checking port " << port << " on " << ip << "..." << std::endl;
    try {
        boost::asio::io_context io_context;
        boost::asio::ip::tcp::socket socket(io_context);
        
        // 使用 make_address 替代 address::from_string
        boost::asio::ip::address addr = boost::asio::ip::make_address(ip);
        boost::asio::ip::tcp::endpoint endpoint(addr, port);
        
        // 使用错误码来跟踪连接状态
        boost::system::error_code ec;
        
        // 设置连接超时
        socket.open(boost::asio::ip::tcp::v4(), ec);
        if (ec) return false;
        
        // 设置非阻塞模式
        socket.non_blocking(true, ec);
        if (ec) return false;
        
        // 尝试连接
        socket.connect(endpoint, ec);
        
        // 如果连接立即成功
        if (!ec) return true;
        
        // 如果连接正在进行中
        if (ec != boost::asio::error::would_block) {
            return false;
        }
        
        // 设置超时
        boost::asio::steady_timer timer(io_context);
        timer.expires_after(std::chrono::milliseconds(500));
        
        // 等待连接完成或超时
        bool connected = false;
        bool timedOut = false;
        
        // 使用 async_wait 替代 async_select
        socket.async_wait(boost::asio::ip::tcp::socket::wait_write,
            [&connected, &ec](const boost::system::error_code& error) {
                if (!error) {
                    // 检查连接是否真的成功
                    boost::system::error_code connect_ec;
                    boost::asio::ip::tcp::socket::endpoint_type ep;
                    connected = !connect_ec;
                } else {
                    ec = error;
                }
            });
        
        // 设置超时处理
        timer.async_wait([&timedOut, &socket](const boost::system::error_code&) {
            timedOut = true;
            socket.cancel();  // 取消正在进行的异步操作
        });
        
        // 运行直到有事件发生
        io_context.run();
        
        // 关闭socket
        socket.close(ec);
        
        // 如果连接成功且没有超时
        return connected && !timedOut;
    }
    catch (const std::exception&) {
        return false;
    }
}

void scan_ports(const std::string& ip, int start_port, int end_port, std::mutex& print_mutex) {
    // 创建线程池，线程数量为硬件支持的并发线程数或自定义数量
    const int thread_count = std::min(std::thread::hardware_concurrency(), 64u);
    std::vector<std::thread> threads;
    std::atomic<int> current_port(start_port);
    
    // 创建多个工作线程
    for (int i = 0; i < thread_count; ++i) {
        threads.emplace_back([&]() {
            while (true) {
                // 获取下一个要扫描的端口
                int port = current_port.fetch_add(1);
                if (port > end_port) {
                    break;  // 所有端口都已分配，退出线程
                }
                
                // 检查端口是否开放
                if (is_port_open(ip, port)) {
                    std::lock_guard<std::mutex> lock(print_mutex);
                    std::cout << "Open port found: " << ip << ":" << port << std::endl;
                }
            }
        });
    }
    
    // 等待所有线程完成
    for (auto& t : threads) {
        if (t.joinable()) {
            t.join();
        }
    }
}